"""RAG（检索增强生成）服务模块

提供文档的解析、向量化、存储和检索功能，支持多种文件格式
实现了多租户隔离，可以根据用户ID和课程ID进行知识隔离
"""

# 导入必要的库和模块
from pathlib import Path
from typing import List, Optional, Tuple
from langchain_community.vectorstores import Chroma  # 向量数据库
from langchain_community.embeddings import SentenceTransformerEmbeddings  # 嵌入模型
from langchain_community.document_loaders import PyPDFLoader, TextLoader  # 文档加载器
from langchain.text_splitter import RecursiveCharacterTextSplitter  # 文本分割器
from langchain.schema import Document  # 文档模型
import io

# 可选导入，避免导入错误
# 检查Pillow库是否可用（用于处理图片）
try:
    from PIL import Image
    PIL_AVAILABLE = True
except ImportError:
    PIL_AVAILABLE = False

# 检查python-docx库是否可用（用于处理Word文档）
try:
    import docx
    DOCX_AVAILABLE = True
except ImportError:
    DOCX_AVAILABLE = False

class RagService:
    """RAG服务类 - 提供文档的解析、向量化、存储和检索功能"""
    
    def __init__(self, persist_dir: str = "app/.chroma", model_name: str = "paraphrase-multilingual-MiniLM-L12-v2"):
        """初始化RAG服务
        
        Args:
            persist_dir: 向量数据库持久化存储目录
            model_name: 用于文本嵌入的模型名称，默认为支持多语言的小型模型
        """
        # 设置持久化目录并确保存在
        self.persist_dir = persist_dir
        Path(self.persist_dir).mkdir(parents=True, exist_ok=True)
        # 初始化嵌入模型，用于将文本转换为向量
        self.embeddings = SentenceTransformerEmbeddings(model_name=model_name)

    def _collection_name(self, user_id: Optional[int], course_id: Optional[int]) -> str:
        """生成集合名称，实现多租户隔离
        
        Args:
            user_id: 用户ID（可选）
            course_id: 课程ID（可选）
        
        Returns:
            str: 基于用户和课程信息的集合名称
        """
        if user_id and course_id:
            return f"user:{user_id}:course:{course_id}"  # 用户+课程级别隔离
        if user_id:
            return f"user:{user_id}"  # 用户级别隔离
        return "global"  # 全局级别

    def _vectorstore(self, collection: str) -> Chroma:
        """获取向量存储实例
        
        Args:
            collection: 集合名称
        
        Returns:
            Chroma: 向量数据库实例
        """
        return Chroma(
            collection_name=collection,
            embedding_function=self.embeddings,
            persist_directory=self.persist_dir
        )

    def ingest_file(self, file_path: str, filename: str, user_id: Optional[int], course_id: Optional[int]) -> int:
        """文档入库：解析 -> 切片 -> 向量化 -> 存储
        
        Args:
            file_path: 文件路径
            filename: 文件名
            user_id: 用户ID（可选）
            course_id: 课程ID（可选）
        
        Returns:
            int: 成功处理并存储的文档块数量
        
        Raises:
            ValueError: 当文件格式不支持或文件内容无法提取时抛出
        """
        # 获取文件扩展名
        ext = Path(filename).suffix.lower()
        
        # 根据文件类型选择不同的处理方式
        if ext == ".pdf":
            docs = self._load_pdf(file_path)
        elif ext == ".txt":
            docs = self._load_txt(file_path)
        elif ext in [".doc", ".docx"]:
            if not DOCX_AVAILABLE:
                raise ValueError("缺少python-docx库，无法处理Word文档")
            docs = self._load_docx(file_path, filename)
        elif ext in [".jpg", ".jpeg", ".png", ".gif", ".bmp", ".webp"]:
            if not PIL_AVAILABLE:
                raise ValueError("缺少Pillow库，无法处理图片文件")
            docs = self._load_image(file_path, filename)
        else:
            raise ValueError(f"不支持的文件格式: {ext}，支持PDF、TXT、DOC、DOCX、图片格式")
        
        if not docs:
            raise ValueError("无法从文件中提取内容")

        # 分块处理（图片不需要分块）
        if ext in [".jpg", ".jpeg", ".png", ".gif", ".bmp", ".webp"]:
            chunks = docs
        else:
            # 使用递归字符分割器，设置块大小和重叠大小
            splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=120)
            chunks = splitter.split_documents(docs)
        
        # 为每个文档块添加元数据
        for chunk in chunks:
            chunk.metadata = chunk.metadata or {}
            chunk.metadata.update({
                "filename": filename,
                "file_type": ext,
                "user_id": user_id,
                "course_id": course_id
            })

        # 存储到向量库
        vs = self._vectorstore(self._collection_name(user_id, course_id))
        vs.add_documents(chunks)  # 添加文档到向量库
        vs.persist()  # 持久化存储
        
        return len(chunks)  # 返回处理的文档块数量
    
    def _load_pdf(self, file_path: str) -> List[Document]:
        """加载PDF文件并解析为文档对象列表
        
        Args:
            file_path: PDF文件路径
        
        Returns:
            List[Document]: 解析后的文档对象列表
        """
        loader = PyPDFLoader(file_path)
        return loader.load()
    
    def _load_txt(self, file_path: str) -> List[Document]:
        """加载TXT文件并解析为文档对象列表
        
        Args:
            file_path: TXT文件路径
        
        Returns:
            List[Document]: 解析后的文档对象列表
        """
        loader = TextLoader(file_path, encoding="utf-8")
        return loader.load()
    
    def _load_docx(self, file_path: str, filename: str) -> List[Document]:
        """加载Word文档并解析为文档对象列表
        
        Args:
            file_path: Word文档路径
            filename: 文件名
        
        Returns:
            List[Document]: 解析后的文档对象列表
        
        Raises:
            ValueError: 当缺少依赖库或文档读取失败时抛出
        """
        if not DOCX_AVAILABLE:
            raise ValueError("缺少python-docx库")
        try:
            doc = docx.Document(file_path)
            text_content = []
            
            # 提取文档中的段落内容
            for paragraph in doc.paragraphs:
                if paragraph.text.strip():
                    text_content.append(paragraph.text.strip())
            
            # 提取文档中的表格内容
            for table in doc.tables:
                for row in table.rows:
                    row_text = []
                    for cell in row.cells:
                        if cell.text.strip():
                            row_text.append(cell.text.strip())
                    if row_text:
                        text_content.append(" | ".join(row_text))
            
            if not text_content:
                return []  # 如果没有提取到内容，返回空列表
                
            # 合并提取的内容
            content = "\n\n".join(text_content)
            return [Document(page_content=content, metadata={"source": filename})]
            
        except Exception as e:
            raise ValueError(f"无法读取Word文档: {str(e)}")
    
    def _load_image(self, file_path: str, filename: str) -> List[Document]:
        """加载图片文件并创建描述性文档对象
        
        Args:
            file_path: 图片文件路径
            filename: 文件名
        
        Returns:
            List[Document]: 包含图片信息的文档对象列表
        
        Raises:
            ValueError: 当缺少依赖库或图片读取失败时抛出
        """
        if not PIL_AVAILABLE:
            raise ValueError("缺少Pillow库")
        try:
            # 获取图片基本信息
            image = Image.open(file_path)
            width, height = image.size
            format_name = image.format or "Unknown"
            
            # 创建图片描述文本
            description = f"""这是一张图片文件：{filename}
格式：{format_name}
尺寸：{width} x {height} 像素
文件类型：图像文件

[图片内容需要通过AI视觉模型分析]
用户可以询问关于此图片的内容、含义或进行分析。
如果用户询问图片内容，系统会调用视觉AI进行分析。"""
            
            # 返回包含图片信息的文档对象
            return [Document(
                page_content=description,
                metadata={
                    "source": filename,
                    "image_width": width,
                    "image_height": height,
                    "image_format": format_name,
                    "file_path": file_path
                }
            )]
            
        except Exception as e:
            raise ValueError(f"无法读取图片文件: {str(e)}")

    def query(self, question: str, user_id: Optional[int], course_id: Optional[int], top_k: int = 4) -> List[Tuple[str, dict]]:
        """检索与问题相关的文档片段
        
        Args:
            question: 用户提问
            user_id: 用户ID（可选）
            course_id: 课程ID（可选）
            top_k: 返回的最相关文档数量，默认为4
        
        Returns:
            List[Tuple[str, dict]]: 包含文档内容和元数据的元组列表
        """
        # 获取对应的向量库
        vs = self._vectorstore(self._collection_name(user_id, course_id))
        # 执行相似度搜索
        results = vs.similarity_search(question, k=top_k)
        # 返回文档内容和元数据
        return [(r.page_content, r.metadata) for r in results]
